package com.chatplus.application.aiprocessor.platform.image.listener;

import cn.bugstack.openai.session.Configuration;
import com.chatplus.application.aiprocessor.platform.image.sd.SdImageProcessor;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.PlusJsonUtils;
import com.chatplus.application.domain.dto.SdJobResult;
import com.chatplus.application.domain.dto.SdParamsDto;
import com.chatplus.application.domain.entity.draw.SdJobEntity;
import com.chatplus.application.service.account.UserProductLogService;
import com.chatplus.application.service.draw.SdJobService;
import okhttp3.*;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.io.IOException;

/**
 * @Description Sd 任务监听
 * 对于延迟任务，需要设置 delayedQueue="true" 并且必须提供 value；否则注解将被忽略。value 是队列名称。设置 deadLetterQueue 可以将任务推送到另一个队列。否则，执行失败时任务会被丢弃。还可以使用 numRetries 字段设置任务重试次数
 * @Author Angus
 * @Date 2024/2/24
 */
@Component
public class SdJobListener {

    private final SdImageProcessor sdImageProcessor;
    private final SdJobService sdJobService;
    private final UserProductLogService userProductLogService;

    public SdJobListener(SdImageProcessor sdImageProcessor, SdJobService sdJobService, UserProductLogService userProductLogService) {
        this.sdImageProcessor = sdImageProcessor;
        this.sdJobService = sdJobService;
        this.userProductLogService = userProductLogService;
    }

    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(SdJobListener.class);
    public void generateImage(SdParamsDto sdParamsDto) {
        OkHttpClient okHttpClient = new OkHttpClient();
        Request request = new Request.Builder()
                .addHeader("Content-Type", Configuration.APPLICATION_JSON)
                .addHeader("Accept", Configuration.APPLICATION_JSON)
                .url("/sdapi/v1/txt2img")
                .post(RequestBody.create(PlusJsonUtils.toJsonString(sdParamsDto), MediaType.parse(Configuration.APPLICATION_JSON)))
                .build();
        Call call = okHttpClient.newCall(request);
        SdJobEntity sdJobEntity = sdJobService.getByTaskId(sdParamsDto.getTaskId());
        try (Response response = call.execute()) {
            if (response.isSuccessful() && response.body() != null) {
                SdJobResult  result  = PlusJsonUtils.parseObject(response.body().string(), SdJobResult.class);
                if (result == null || result.getImages() == null || result.getImages().length == 0) {
                    LOGGER.message("SD图片生成失败").context("response",response).error();
                    sdJobEntity.setProgress(-1);
                    return;
                }
                String url = sdImageProcessor.saveToOss(null, result.getImages()[0]);
                if (StringUtils.isNotEmpty(url)) {
                    sdJobEntity.setImgUrl(url);
                    sdJobEntity.setProgress(100);
                    userProductLogService.reduceImgCalls(sdJobEntity.getUserId(), 1);
                }
            } else {
                LOGGER.message("SD图片请求失败").context("response",response).error();
            }
        } catch (IOException e) {
            sdJobEntity.setProgress(-1);
            LOGGER.message("SD图片请求失败").exception(e).error();
        }finally {
            sdJobService.updateById(sdJobEntity);
        }
    }
}
